#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
安全中间件

这个模块包含安全相关的中间件：
1. CORS（跨域资源共享）处理
2. 安全头部设置
3. IP白名单/黑名单
4. 请求大小限制
5. 安全事件记录

提供全面的安全防护功能
支持自定义安全策略

作者: FastAPI 学习项目
创建时间: 2024
"""

import re
import ipaddress
from typing import Callable, List, Optional, Set, Dict, Any
from fastapi import Request, Response, HTTPException
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.middleware.cors import CORSMiddleware
from loguru import logger
from datetime import datetime, timedelta
from collections import defaultdict, deque
import threading
import hashlib
import secrets


class SecurityConfig:
    """
    安全配置类
    
    定义各种安全策略和配置参数。
    """
    
    def __init__(
        self,
        *,
        # CORS 配置
        allow_origins: List[str] = None,
        allow_credentials: bool = True,
        allow_methods: List[str] = None,
        allow_headers: List[str] = None,
        
        # 安全头部配置
        enable_security_headers: bool = True,
        content_security_policy: str = None,
        
        # IP 访问控制
        ip_whitelist: List[str] = None,
        ip_blacklist: List[str] = None,
        
        # 请求限制
        max_request_size: int = 10 * 1024 * 1024,  # 10MB
        max_header_size: int = 8192,  # 8KB
        
        # 安全检查
        enable_sql_injection_check: bool = True,
        enable_xss_check: bool = True,
        enable_path_traversal_check: bool = True,
        
        # 其他配置
        trusted_hosts: List[str] = None,
        block_suspicious_user_agents: bool = True
    ):
        """
        初始化安全配置
        
        Args:
            allow_origins: 允许的源域名列表
            allow_credentials: 是否允许凭据
            allow_methods: 允许的HTTP方法列表
            allow_headers: 允许的请求头列表
            enable_security_headers: 是否启用安全头部
            content_security_policy: 内容安全策略
            ip_whitelist: IP白名单
            ip_blacklist: IP黑名单
            max_request_size: 最大请求大小
            max_header_size: 最大请求头大小
            enable_sql_injection_check: 是否启用SQL注入检查
            enable_xss_check: 是否启用XSS检查
            enable_path_traversal_check: 是否启用路径遍历检查
            trusted_hosts: 受信任的主机列表
            block_suspicious_user_agents: 是否阻止可疑的User-Agent
        """
        # CORS 配置
        self.allow_origins = allow_origins or ["*"]
        self.allow_credentials = allow_credentials
        self.allow_methods = allow_methods or ["GET", "POST", "PUT", "DELETE", "OPTIONS"]
        self.allow_headers = allow_headers or ["*"]
        
        # 安全头部配置
        self.enable_security_headers = enable_security_headers
        self.content_security_policy = content_security_policy or (
            "default-src 'self'; "
            "script-src 'self' 'unsafe-inline' 'unsafe-eval' https://cdn.jsdelivr.net https://unpkg.com; "
            "style-src 'self' 'unsafe-inline' https://cdn.jsdelivr.net https://unpkg.com; "
            "img-src 'self' data: https:; "
            "font-src 'self' data: https://cdn.jsdelivr.net https://unpkg.com; "
            "connect-src 'self';"
        )
        
        # IP 访问控制
        self.ip_whitelist = self._parse_ip_list(ip_whitelist or [])
        self.ip_blacklist = self._parse_ip_list(ip_blacklist or [])
        
        # 请求限制
        self.max_request_size = max_request_size
        self.max_header_size = max_header_size
        
        # 安全检查
        self.enable_sql_injection_check = enable_sql_injection_check
        self.enable_xss_check = enable_xss_check
        self.enable_path_traversal_check = enable_path_traversal_check
        
        # 其他配置
        self.trusted_hosts = trusted_hosts or []
        self.block_suspicious_user_agents = block_suspicious_user_agents
        
        # 可疑User-Agent模式
        self.suspicious_user_agents = [
            r'.*bot.*',
            r'.*crawler.*',
            r'.*spider.*',
            r'.*scraper.*',
            r'.*scanner.*',
            r'.*sqlmap.*',
            r'.*nikto.*',
            r'.*nmap.*'
        ]
    
    def _parse_ip_list(self, ip_list: List[str]) -> List:
        """
        解析IP列表，支持单个IP和CIDR网段
        
        Args:
            ip_list: IP字符串列表
            
        Returns:
            List: 解析后的IP网络对象列表
        """
        parsed_ips = []
        for ip_str in ip_list:
            try:
                if '/' in ip_str:
                    # CIDR 网段
                    parsed_ips.append(ipaddress.ip_network(ip_str, strict=False))
                else:
                    # 单个IP
                    parsed_ips.append(ipaddress.ip_network(f"{ip_str}/32", strict=False))
            except ValueError as e:
                logger.warning(f"无效的IP地址格式: {ip_str}, 错误: {e}")
        return parsed_ips


class SecurityEventTracker:
    """
    安全事件跟踪器
    
    记录和分析安全相关事件。
    """
    
    def __init__(self, max_events: int = 1000):
        """
        初始化安全事件跟踪器
        
        Args:
            max_events: 最大事件记录数量
        """
        self.max_events = max_events
        self.security_events = deque(maxlen=max_events)
        self.ip_attempts = defaultdict(lambda: {'count': 0, 'last_attempt': None})
        self.blocked_ips = set()
        self._lock = threading.Lock()
    
    def record_event(
        self,
        event_type: str,
        client_ip: str,
        details: Dict[str, Any],
        severity: str = "medium"
    ):
        """
        记录安全事件
        
        Args:
            event_type: 事件类型
            client_ip: 客户端IP
            details: 事件详情
            severity: 严重程度 (low, medium, high, critical)
        """
        with self._lock:
            event = {
                'timestamp': datetime.utcnow(),
                'event_type': event_type,
                'client_ip': client_ip,
                'details': details,
                'severity': severity
            }
            
            self.security_events.append(event)
            
            # 更新IP尝试计数
            self.ip_attempts[client_ip]['count'] += 1
            self.ip_attempts[client_ip]['last_attempt'] = datetime.utcnow()
            
            # 记录日志
            logger.warning(
                f"安全事件: {event_type}",
                extra={
                    "event_type": "security_event",
                    "security_event_type": event_type,
                    "client_ip": client_ip,
                    "severity": severity,
                    "details": details,
                    "timestamp": event['timestamp'].isoformat()
                }
            )
    
    def should_block_ip(self, client_ip: str, threshold: int = 10) -> bool:
        """
        判断是否应该阻止某个IP
        
        Args:
            client_ip: 客户端IP
            threshold: 阻止阈值
            
        Returns:
            bool: 是否应该阻止
        """
        with self._lock:
            if client_ip in self.blocked_ips:
                return True
            
            attempts = self.ip_attempts.get(client_ip, {'count': 0})
            if attempts['count'] >= threshold:
                self.blocked_ips.add(client_ip)
                logger.critical(f"IP {client_ip} 已被自动阻止，尝试次数: {attempts['count']}")
                return True
            
            return False
    
    def get_recent_events(self, limit: int = 50) -> List[Dict]:
        """
        获取最近的安全事件
        
        Args:
            limit: 返回数量限制
            
        Returns:
            List[Dict]: 安全事件列表
        """
        with self._lock:
            return list(self.security_events)[-limit:]
    
    def get_ip_statistics(self) -> Dict:
        """
        获取IP统计信息
        
        Returns:
            Dict: IP统计信息
        """
        with self._lock:
            return {
                'total_ips': len(self.ip_attempts),
                'blocked_ips': len(self.blocked_ips),
                'top_offenders': sorted(
                    self.ip_attempts.items(),
                    key=lambda x: x[1]['count'],
                    reverse=True
                )[:10]
            }


class SecurityMiddleware(BaseHTTPMiddleware):
    """
    安全中间件
    
    提供全面的安全防护功能。
    """
    
    def __init__(self, app, config: SecurityConfig = None):
        """
        初始化安全中间件
        
        Args:
            app: FastAPI 应用实例
            config: 安全配置
        """
        super().__init__(app)
        self.config = config or SecurityConfig()
        self.event_tracker = SecurityEventTracker()
        
        # 编译正则表达式模式
        self._compile_security_patterns()
    
    def _compile_security_patterns(self):
        """
        编译安全检查的正则表达式模式
        """
        # SQL注入模式 - 更精确的检测
        self.sql_injection_patterns = [
            # 检测SQL注入关键字组合
            re.compile(r"\b(union\s+select|select\s+.*\s+from|insert\s+into|delete\s+from|update\s+.*\s+set|drop\s+table|create\s+table)", re.IGNORECASE),
            # 检测SQL注释和终止符组合
            re.compile(r"(--|/\*|\*/|;\s*--)", re.IGNORECASE),
            # 检测常见SQL注入payload
            re.compile(r"('\s*(or|and)\s*'|'\s*(or|and)\s*\d+\s*=\s*\d+|'\s*;\s*drop|'\s*;\s*delete)", re.IGNORECASE),
            # 检测十六进制和字符编码注入
            re.compile(r"(0x[0-9a-f]+|char\(|ascii\(|substring\()", re.IGNORECASE)
        ]
        
        # XSS模式
        self.xss_patterns = [
            re.compile(r"<script[^>]*>.*?</script>", re.IGNORECASE | re.DOTALL),
            re.compile(r"javascript:", re.IGNORECASE),
            re.compile(r"on\w+\s*=", re.IGNORECASE),
            re.compile(r"<iframe[^>]*>", re.IGNORECASE)
        ]
        
        # 路径遍历模式
        self.path_traversal_patterns = [
            re.compile(r"\.\./"),
            re.compile(r"\.\.\\"),
            re.compile(r"%2e%2e%2f", re.IGNORECASE),
            re.compile(r"%2e%2e%5c", re.IGNORECASE)
        ]
        
        # 可疑User-Agent模式
        self.suspicious_ua_patterns = [
            re.compile(pattern, re.IGNORECASE) 
            for pattern in self.config.suspicious_user_agents
        ]
    
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        """
        处理请求并执行安全检查
        
        Args:
            request: HTTP请求对象
            call_next: 下一个中间件或路由处理器
            
        Returns:
            Response: HTTP响应对象
        """
        client_ip = self._get_client_ip(request)
        
        try:
            # 执行安全检查
            await self._perform_security_checks(request, client_ip)
            
            # 处理请求
            response = await call_next(request)
            
            # 添加安全头部
            if self.config.enable_security_headers:
                self._add_security_headers(response)
            
            return response
            
        except HTTPException as e:
            # 记录安全事件
            self.event_tracker.record_event(
                event_type="security_block",
                client_ip=client_ip,
                details={
                    "reason": str(e.detail),
                    "status_code": e.status_code,
                    "path": request.url.path,
                    "method": request.method
                },
                severity="high"
            )
            raise
    
    def _get_client_ip(self, request: Request) -> str:
        """
        获取客户端真实IP地址
        
        Args:
            request: HTTP请求对象
            
        Returns:
            str: 客户端IP地址
        """
        # 检查代理头部
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            return forwarded_for.split(",")[0].strip()
        
        real_ip = request.headers.get("X-Real-IP")
        if real_ip:
            return real_ip
        
        # 使用客户端IP
        return request.client.host if request.client else "unknown"
    
    async def _perform_security_checks(self, request: Request, client_ip: str):
        """
        执行安全检查
        
        Args:
            request: HTTP请求对象
            client_ip: 客户端IP地址
        """
        # IP访问控制检查
        await self._check_ip_access(client_ip)
        
        # 请求大小检查
        await self._check_request_size(request)
        
        # 主机检查
        await self._check_trusted_hosts(request)
        
        # User-Agent检查
        await self._check_user_agent(request, client_ip)
        
        # 安全模式检查
        await self._check_security_patterns(request, client_ip)
    
    async def _check_ip_access(self, client_ip: str):
        """
        检查IP访问控制
        
        Args:
            client_ip: 客户端IP地址
        """
        try:
            client_ip_obj = ipaddress.ip_address(client_ip)
        except ValueError:
            # 无效IP地址
            raise HTTPException(status_code=403, detail="无效的IP地址")
        
        # 检查是否被自动阻止
        if self.event_tracker.should_block_ip(client_ip):
            raise HTTPException(status_code=403, detail="IP地址已被阻止")
        
        # 检查黑名单
        for blocked_network in self.config.ip_blacklist:
            if client_ip_obj in blocked_network:
                raise HTTPException(status_code=403, detail="IP地址在黑名单中")
        
        # 检查白名单（如果配置了白名单）
        if self.config.ip_whitelist:
            allowed = False
            for allowed_network in self.config.ip_whitelist:
                if client_ip_obj in allowed_network:
                    allowed = True
                    break
            
            if not allowed:
                raise HTTPException(status_code=403, detail="IP地址不在白名单中")
    
    async def _check_request_size(self, request: Request):
        """
        检查请求大小
        
        Args:
            request: HTTP请求对象
        """
        # 检查Content-Length
        content_length = request.headers.get("content-length")
        if content_length:
            try:
                size = int(content_length)
                if size > self.config.max_request_size:
                    raise HTTPException(
                        status_code=413, 
                        detail=f"请求体过大: {size} > {self.config.max_request_size}"
                    )
            except ValueError:
                raise HTTPException(status_code=400, detail="无效的Content-Length")
        
        # 检查请求头大小
        headers_size = sum(len(k) + len(v) for k, v in request.headers.items())
        if headers_size > self.config.max_header_size:
            raise HTTPException(
                status_code=431,
                detail=f"请求头过大: {headers_size} > {self.config.max_header_size}"
            )
    
    async def _check_trusted_hosts(self, request: Request):
        """
        检查受信任的主机
        
        Args:
            request: HTTP请求对象
        """
        if not self.config.trusted_hosts:
            return
        
        host = request.headers.get("host", "")
        if host not in self.config.trusted_hosts:
            raise HTTPException(status_code=400, detail="不受信任的主机")
    
    async def _check_user_agent(self, request: Request, client_ip: str):
        """
        检查User-Agent
        
        Args:
            request: HTTP请求对象
            client_ip: 客户端IP地址
        """
        if not self.config.block_suspicious_user_agents:
            return
        
        user_agent = request.headers.get("user-agent", "")
        
        # 检查是否为空或可疑的User-Agent
        if not user_agent:
            self.event_tracker.record_event(
                event_type="suspicious_user_agent",
                client_ip=client_ip,
                details={"user_agent": "empty", "path": request.url.path},
                severity="medium"
            )
            return
        
        # 检查可疑模式
        for pattern in self.suspicious_ua_patterns:
            if pattern.search(user_agent):
                self.event_tracker.record_event(
                    event_type="suspicious_user_agent",
                    client_ip=client_ip,
                    details={
                        "user_agent": user_agent,
                        "pattern": pattern.pattern,
                        "path": request.url.path
                    },
                    severity="high"
                )
                raise HTTPException(status_code=403, detail="可疑的User-Agent")
    
    async def _check_security_patterns(self, request: Request, client_ip: str):
        """
        检查安全模式（SQL注入、XSS、路径遍历等）
        
        Args:
            request: HTTP请求对象
            client_ip: 客户端IP地址
        """
        # 白名单路径，跳过安全检查
        whitelist_paths = [
            "/docs", "/redoc", "/openapi.json", "/favicon.ico",
            "/health", "/metrics", "/static"
        ]
        
        # 检查是否为白名单路径
        for whitelist_path in whitelist_paths:
            if request.url.path.startswith(whitelist_path):
                return
        
        # 获取要检查的内容
        check_content = []
        
        # 检查URL路径
        check_content.append(("path", request.url.path))
        
        # 检查查询参数
        for key, value in request.query_params.items():
            check_content.append(("query_param", f"{key}={value}"))
        
        # 检查请求头（排除一些标准头部）
        excluded_headers = {"authorization", "cookie", "user-agent", "host", "accept", "accept-encoding", "accept-language", "connection", "content-type"}
        for key, value in request.headers.items():
            if key.lower() not in excluded_headers:
                check_content.append(("header", f"{key}: {value}"))
        
        # 执行模式检查
        for content_type, content in check_content:
            # SQL注入检查
            if self.config.enable_sql_injection_check:
                for pattern in self.sql_injection_patterns:
                    if pattern.search(content):
                        self.event_tracker.record_event(
                            event_type="sql_injection_attempt",
                            client_ip=client_ip,
                            details={
                                "content_type": content_type,
                                "content": content,
                                "pattern": pattern.pattern
                            },
                            severity="critical"
                        )
                        raise HTTPException(status_code=403, detail="检测到SQL注入尝试")
            
            # XSS检查
            if self.config.enable_xss_check:
                for pattern in self.xss_patterns:
                    if pattern.search(content):
                        self.event_tracker.record_event(
                            event_type="xss_attempt",
                            client_ip=client_ip,
                            details={
                                "content_type": content_type,
                                "content": content,
                                "pattern": pattern.pattern
                            },
                            severity="high"
                        )
                        raise HTTPException(status_code=403, detail="检测到XSS尝试")
            
            # 路径遍历检查
            if self.config.enable_path_traversal_check:
                for pattern in self.path_traversal_patterns:
                    if pattern.search(content):
                        self.event_tracker.record_event(
                            event_type="path_traversal_attempt",
                            client_ip=client_ip,
                            details={
                                "content_type": content_type,
                                "content": content,
                                "pattern": pattern.pattern
                            },
                            severity="high"
                        )
                        raise HTTPException(status_code=403, detail="检测到路径遍历尝试")
    
    def _add_security_headers(self, response: Response):
        """
        添加安全头部
        
        Args:
            response: HTTP响应对象
        """
        # 内容安全策略
        response.headers["Content-Security-Policy"] = self.config.content_security_policy
        
        # 防止点击劫持
        response.headers["X-Frame-Options"] = "DENY"
        
        # 防止MIME类型嗅探
        response.headers["X-Content-Type-Options"] = "nosniff"
        
        # XSS保护
        response.headers["X-XSS-Protection"] = "1; mode=block"
        
        # 强制HTTPS
        response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains"
        
        # 引用策略
        response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
        
        # 权限策略
        response.headers["Permissions-Policy"] = (
            "geolocation=(), microphone=(), camera=()"
        )
        
        # 服务器信息隐藏
        response.headers["Server"] = "FastAPI"
        
        # 缓存控制（对于敏感页面）
        if "login" in str(response.headers.get("location", "")).lower():
            response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
            response.headers["Pragma"] = "no-cache"
            response.headers["Expires"] = "0"
    
    def get_security_report(self) -> Dict:
        """
        获取安全报告
        
        Returns:
            Dict: 安全报告数据
        """
        return {
            "recent_events": self.event_tracker.get_recent_events(),
            "ip_statistics": self.event_tracker.get_ip_statistics(),
            "config_summary": {
                "ip_whitelist_count": len(self.config.ip_whitelist),
                "ip_blacklist_count": len(self.config.ip_blacklist),
                "security_checks_enabled": {
                    "sql_injection": self.config.enable_sql_injection_check,
                    "xss": self.config.enable_xss_check,
                    "path_traversal": self.config.enable_path_traversal_check
                },
                "max_request_size": self.config.max_request_size,
                "trusted_hosts_count": len(self.config.trusted_hosts)
            }
        }